from collections import Vector
from math import f64, max, min

struct DoseStatRet:
    binned_dose_values: Vector[f64]
    accumulated_relative_volumes: Vector[f64]

def Accumulate(
    voxel_volume: Vector[f64],
    voxel_dose: Vector[f64],
    do_exact_accumulation: Bool = False
) -> DoseStatRet:
    # 过滤体积为0的体素
    dose_volume_pairs = Vector[(f64, f64)]()
    for i in range(voxel_volume.size()):
        if voxel_volume[i] > 0.0:
            dose = max(0.0, voxel_dose[i])
            dose_volume_pairs.append((dose, voxel_volume[i]))
    
    num_pstv = dose_volume_pairs.size()
    
    # 处理空数据
    if num_pstv == 0:
        return DoseStatRet(
            binned_dose_values=[0.0, 0.0],
            accumulated_relative_volumes=[1.0, 0.0]
        )
    
    # 选择计算方法
    EXACT_THRESHOLD = 5000
    use_exact_method = do_exact_accumulation or (num_pstv < EXACT_THRESHOLD)
    
    if use_exact_method:
        # 精确计算：直接排序
        dose_volume_pairs.sort(key=lambda pair: pair[0])
        
        # 准备结果
        binned_dose_values = Vector[f64](num_pstv + 2)
        accumulated_relative_volumes = Vector[f64](num_pstv + 2)
        
        # 初始化向量
        for i in range(binned_dose_values.size()):
            binned_dose_values[i] = 0.0
        for i in range(accumulated_relative_volumes.size()):
            accumulated_relative_volumes[i] = 0.0
        
        # 填充排序后的结果
        for i in range(num_pstv):
            binned_dose_values[i + 1] = dose_volume_pairs[i][0]
            accumulated_relative_volumes[i + 1] = dose_volume_pairs[i][1]
        
        # 填充哨兵
        binned_dose_values[0] = binned_dose_values[1]
        binned_dose_values[-1] = binned_dose_values[num_pstv]
        accumulated_relative_volumes[0] = 1.0
        accumulated_relative_volumes[-1] = 0.0
        
        # 归一化
        total_vol = 0.0
        for i in range(accumulated_relative_volumes.size()):
            total_vol += accumulated_relative_volumes[i]
        
        if total_vol > 0.0:
            for i in range(accumulated_relative_volumes.size()):
                accumulated_relative_volumes[i] /= total_vol
        
        return DoseStatRet(
            binned_dose_values=binned_dose_values,
            accumulated_relative_volumes=accumulated_relative_volumes
        )
    else:
        # 粗略计算
        doses = Vector[f64]()
        total_volume = 0.0
        for (dose, volume) in dose_volume_pairs:
            doses.append(dose)
            total_volume += volume
        
        # 计算剂量范围
        min_dose = max(0.0, doses.min())
        max_dose = doses.max()
        
        dose_range = max_dose - min_dose
        
        # 处理剂量范围过小的情况
        if dose_range <= 1e-6:
            raise ValueError("Dose range too small")
        
        # 动态计算分箱数
        MIN_BINS = 2
        MAX_BINS = 1001
        dose_min_width = 1e-6
        theoretical_bins = dose_range / dose_min_width + 1.0
        num_bins = max(MIN_BINS, min(MAX_BINS, int(theoretical_bins)))
        
        # 准备分箱结果
        binned_dose_values = Vector[f64](num_bins + 2)
        accumulated_relative_volumes = Vector[f64](num_bins + 2)
        
        # 初始化向量
        for i in range(binned_dose_values.size()):
            binned_dose_values[i] = 0.0
        for i in range(accumulated_relative_volumes.size()):
            accumulated_relative_volumes[i] = 0.0
        
        # 填充分箱剂量值
        bin_width = dose_range / (num_bins - 1)
        for j in range(num_bins):
            binned_dose_values[j + 1] = min_dose + j * bin_width
        
        # 统计每个分箱的体积
        for (dose, volume) in dose_volume_pairs:
            if dose <= min_dose:
                accumulated_relative_volumes[1] += volume
            elif dose >= max_dose:
                # 不处理，因为已经超出范围
                pass
            else:
                bin_index = int((dose - min_dose) / bin_width + 0.5)
                accumulated_relative_volumes[bin_index + 1] += volume
        
        # 计算累积相对体积
        total_vol = 0.0
        for i in range(num_bins, 0, -1):
            vol = accumulated_relative_volumes[i]
            accumulated_relative_volumes[i] = total_vol + 0.5 * vol
            total_vol += vol
        
        # 填充哨兵
        accumulated_relative_volumes[0] = total_vol
        accumulated_relative_volumes[-1] = 0.0
        binned_dose_values[0] = binned_dose_values[1]
        binned_dose_values[-1] = binned_dose_values[num_bins]
        
        # 归一化
        if total_vol > 0.0:
            for i in range(accumulated_relative_volumes.size()):
                accumulated_relative_volumes[i] /= total_vol
        
        return DoseStatRet(
            binned_dose_values=binned_dose_values,
            accumulated_relative_volumes=accumulated_relative_volumes
        )

def GetMetric(
    dose_at_volume: f64,
    binned_dose_values: Vector[f64],
    accumulated_relative_volumes: Vector[f64]
) -> f64:
    # 参数校验
    if accumulated_relative_volumes.size() == 0 or binned_dose_values.size() == 0:
        print("Warning: Empty input vectors provided to GetMetric")
        return 0.0
    
    # 边界检查
    if dose_at_volume >= accumulated_relative_volumes[0]:
        return binned_dose_values[0]
    if dose_at_volume <= accumulated_relative_volumes[-1]:
        return binned_dose_values[-1]
    
    # 二分查找
    low = 0
    high = accumulated_relative_volumes.size() - 1
    while low <= high:
        mid = (low + high) // 2
        if accumulated_relative_volumes[mid] == dose_at_volume:
            # 找到匹配的值
            # 计算匹配值的平均剂量
            start = mid
            end = mid
            while start > 0 and accumulated_relative_volumes[start - 1] == dose_at_volume:
                start -= 1
            while end < accumulated_relative_volumes.size() - 1 and accumulated_relative_volumes[end + 1] == dose_at_volume:
                end += 1
            total_dose = 0.0
            count = 0
            for i in range(start, end + 1):
                total_dose += binned_dose_values[i]
                count += 1
            return total_dose / count
        elif accumulated_relative_volumes[mid] > dose_at_volume:
            low = mid + 1
        else:
            high = mid - 1
    
    # 如果没有精确匹配，进行线性插值
    idx = low
    if idx == 0:
        return binned_dose_values[0]
    if idx >= accumulated_relative_volumes.size():
        return binned_dose_values[-1]
    
    idx1 = idx - 1
    idx2 = idx
    
    vol_diff = accumulated_relative_volumes[idx2] - accumulated_relative_volumes[idx1]
    if abs(vol_diff) <= 1e-10:
        return 0.5 * (binned_dose_values[idx2] + binned_dose_values[idx1])
    
    dose_diff = binned_dose_values[idx2] - binned_dose_values[idx1]
    metric = binned_dose_values[idx1] + (dose_at_volume - accumulated_relative_volumes[idx1]) * dose_diff / vol_diff
    return metric

def GetDoseStatRoiInterp(
    dose_type: String,
    binned_dose_values: Vector[f64],
    accumulated_relative_volumes: Vector[f64],
    voxel_volume: Vector[f64],
    voxel_dose: Vector[f64],
    dose_at_volume: f64 = 0.0
) -> f64:
    # 检查数组大小
    if voxel_volume.size() != voxel_dose.size():
        raise ValueError("voxelVolume and voxelDose arrays have different sizes")
    
    # 计算平均剂量
    if dose_type == "average":
        total_volume = 0.0
        weighted_dose = 0.0
        for i in range(voxel_volume.size()):
            total_volume += voxel_volume[i]
            weighted_dose += voxel_dose[i] * voxel_volume[i]
        
        if total_volume <= 0.0:
            raise ValueError("Prescription ROI volume is zero. Cannot scale to average dose.")
        
        return weighted_dose / total_volume
    
    # 提取目标体积百分比
    target_dose_at_volume: f64
    if dose_type == "D_98":
        target_dose_at_volume = 0.98
    elif dose_type == "D_95":
        target_dose_at_volume = 0.95
    elif dose_type == "D_50":
        target_dose_at_volume = 0.5
    elif dose_type == "D_2":
        target_dose_at_volume = 0.02
    elif dose_type == "DoseAtVolume":
        target_dose_at_volume = dose_at_volume
    else:
        raise ValueError("Prescription type not supported. Only 'average', 'D_98', 'D_95', 'D_50', 'D_2', 'DoseAtVolume' are supported")
    
    # 调用GetMetric
    return GetMetric(target_dose_at_volume, binned_dose_values, accumulated_relative_volumes)